# Define statistical models

# Load packages
```{r}
library(rethinking)
```

# Load data
```{r}
# Load data from module_1
load("./data/d_sol_cw_w_in.RData")

# Load simulation data from module_2
# load("./data/d_sim_in.RData")
```


## Choose between real-world or simulated data
```{r}
# d_in <- d_sim_m1 # simulated data of assumption m1
# d_in <- d_sim_m2 # simulated data of assumption m2
# d_in <- d_sim_m3 # simulated data of assumption m3
d_in <- d_sol_cw # real world (tm) data
```

## 1. Define model m1: (T_S <-) Sol -> C_MW
```{r}
# Define the input data
dlist <- list(
    C_MW = (d_in$C_MW),
    Sol_std = (d_in$Sol_std),
    T_S_std = (d_in$T_S_std)
)

model_m1 <- ulam(log_lik=TRUE,
    alist(
    # Sol -> C_MW model
    C_MW ~ dnorm(mu, sigma),
    mu <- a_C + b_C_S * Sol_std,
    a_C ~ dnorm(0, 1),
    b_C_S ~ dnorm(0.5, 0.5),
    sigma ~ dexp(1),

    # Sol -> T_S model
    T_S_std ~ dnorm(nu, tau),
    nu <- a_T + b_T_S * Sol_std,
    a_T ~ dnorm(0.5, 1),
    b_T_S ~ dnorm(0.5, 0.5),
    tau ~ dexp(1)
    ), data = dlist, chains = 2, cores = 2
)

```

## 2. Define model m2: Sol -> T_S -> C_MW
```{r}
# Define the input data
dlist <- list(
    C_MW = (d_in$C_MW),
    Sol_std = (d_in$Sol_std),
    T_S_std = (d_in$T_S_std)
)

model_m2 <- ulam(log_lik=TRUE,
    alist(
    # T_S -> C_MW model 
    C_MW ~ dnorm(mu, sigma),
    mu <- a_C + b_C_T * T_S_std,
    a_C ~ dnorm(0, 1),
    b_C_T ~ dnorm(0.5, 0.5),
    sigma ~ dexp(1),

    # Sol -> T_S model
    T_S_std ~ dnorm(nu, tau),
    nu <- a_T + b_T_S * Sol_std,
    a_T ~ dnorm(0.5, 1),
    b_T_S ~ dnorm(0.5, 0.5),
    tau ~ dexp(1)
    ), data = dlist, chains = 2, cores = 2
)

```

## 3. Define model m3: T_S -> C_MW <- Sol with interaction T_S -> Sol
```{r}
# Define the input data
dlist <- list(
    C_MW = (d_in$C_MW),
    Sol_std = (d_in$Sol_std),
    T_S_std = (d_in$T_S_std)
)

model_m3 <- ulam(log_lik=TRUE,
    alist(
    # Sol -> C_MW <- T_S model
    C_MW ~ dnorm(mu, sigma),
    mu <- a_C + b_C_S * Sol_std + b_C_T * T_S_std,
    a_C ~ dnorm(0, 1),
    b_C_S ~ dnorm(0.5, 0.5),
    b_C_T ~ dnorm(0.5, 0.5),
    sigma ~ dexp(1),

    # Sol -> T_S model
    T_S_std ~ dnorm(nu, tau),
    nu <- a_T + b_T_S * Sol_std,
    a_T ~ dnorm(0.5, 1),
    b_T_S ~ dnorm(0.5, 0.5),
    tau ~ dexp(1)
    ), data = dlist, chains = 2, cores = 2
)

```

## 5. Define model m4: T_S -> C_MW <- Sol with interaction Sol -> T_S and T_S <- W -> C_MW and external factor r_W -> C_MW influenced by Sol, T_S and W
```{r}
dlist <- list(
    C_MW = (d_in$C_MW),
    W_std = standardize(d_in$W),
    Sol_std = (d_in$Sol_std),
    T_S_std = (d_in$T_S_std),
    r_W_std = (d_in$r_W_std)
)

model_m4 <- ulam(log_lik=TRUE,
    alist(
    # Sol -> C_MW <- T_S model and r_w -> C_MW <- W
    C_MW ~ dnorm(mu, sigma),
    mu <- a_C + b_C_S * Sol_std + b_C_T * T_S_std + b_C_W * W_std + b_C_r * r_W_std,
    a_C ~ dnorm(0, 1),
    b_C_S ~ dnorm(0.5, 0.5),
    b_C_T ~ dnorm(0.5, 0.5),
    b_C_W ~ dnorm(0.5, 0.5),
    b_C_r ~ dnorm(0.5, 0.5),
    sigma ~ dexp(1),

    # Sol -> T_S model <- W
    T_S_std ~ dnorm(nu, tau),
    nu <- a_T + b_T_S * Sol_std + b_T_W * W_std,
    a_T ~ dnorm(0.5, 1),
    b_T_S ~ dnorm(0.5, 0.5),
    b_T_W ~ dnorm(0.5, 0.5),
    tau ~ dexp(1),

    # T_S -> r_w model
    r_W_std ~ dnorm(xi, ypsilon),
    xi <- a_r + b_r_T * T_S_std + b_r_W * W_std + b_r_S * Sol_std,
    a_r ~ dnorm(0, 1),
    b_r_T ~ dnorm(0.5, 0.5),
    b_r_W ~ dnorm(0.5, 0.5),
    b_r_S ~ dnorm(0.5, 0.5),
    ypsilon ~ dexp(1)

    ), data = dlist, chains = 2, cores = 2
)

precis(model_m4)
```

## 6. Define model m5: T_S -> C_MW <- Sol with interaction W <- Sol -> T_S and T_S <- W -> C_MW and external factor r_W -> C_MW influenced by Sol, T_S and W
```{r}
dlist <- list(
    C_MW = (d_in$C_MW),
    W_std = standardize(d_in$W),
    Sol_std = (d_in$Sol_std),
    T_S_std = (d_in$T_S_std),
    r_W_std = (d_in$r_W_std)
)

model_m5 <- ulam(log_lik=TRUE,
    alist(
    # Sol -> C_MW <- T_S model and C_MW <- r_W
    C_MW ~ dnorm(mu, sigma),
    mu <- a_C + b_C_S * Sol_std + b_C_T * T_S_std + b_C_W * W_std + b_C_r * r_W_std,
    a_C ~ dnorm(0, 1),
    b_C_S ~ dnorm(0.5, 0.5),
    b_C_T ~ dnorm(0.5, 0.5),
    b_C_W ~ dnorm(0.5, 0.5),
    b_C_r ~ dnorm(0.5, 0.5),
    sigma ~ dexp(1),

    # Sol -> T_S model <- W
    T_S_std ~ dnorm(nu, tau),
    nu <- a_T + b_T_S * Sol_std + b_T_W * W_std,
    a_T ~ dnorm(0.5, 1),
    b_T_S ~ dnorm(0.5, 0.5),
    b_T_W ~ dnorm(0.5, 0.5),
    tau ~ dexp(1),

    # T_S -> r_w model
    r_W_std ~ dnorm(xi, ypsilon),
    xi <- a_r + b_r_T * T_S_std + b_r_W * W_std + b_r_S * Sol_std,
    a_r ~ dnorm(0, 1),
    b_r_T ~ dnorm(0.5, 0.5),
    b_r_W ~ dnorm(0.5, 0.5),
    b_r_S ~ dnorm(0.5, 0.5),
    ypsilon ~ dexp(1),

    # Sol -> W model
    W_std ~ dnorm(psi, phi),
    psi <- a_W + b_W_S * Sol_std,
    a_W ~ dnorm(0, 1),
    b_W_S ~ dnorm(0.5, 0.5),
    phi ~ dexp(1)

    ), data = dlist, chains = 2, cores = 2
)

precis(model_m5)
```

## Saving the models so we do not have to re-run them all the time.
```{r}
save(model_m1, model_m2, model_m3, model_m4, model_m5, file = "./data/model_out_final_real.RDATA")

precis(model_m1, prob=0.9, digits = 3)
precis(model_m2, prob=0.9, digits = 3)
precis(model_m3, prob=0.9, digits = 3)
precis(model_m4, prob=0.9, digits = 3)
precis(model_m5, prob=0.9, digits = 3)

windows();
compare(model_m1, model_m2, model_m3, model_m4, model_m5)
``` 

## Exporting the data from the models for plotting in seaborn
```{r}
data_m1 <- extract.samples(model_m1)
data_m2 <- extract.samples(model_m2)
data_m3 <- extract.samples(model_m3)
data_m4 <- extract.samples(model_m4)
data_m5 <- extract.samples(model_m5)

save(data_m1, data_m2, data_m3, data_m4, data_m5, file = "./data/modeldata_out_final_real.RData")
```

## Calculate correlations
```{r}
vcov(model_m3)
```

## Run all above
```{r}

``` 
